package com.yc.grpc.interceptor;

import com.google.common.annotations.VisibleForTesting;
import io.grpc.*;

import java.util.logging.Logger;

/**
 * @Author: ZhouJianLong
 * @Date: 2024/07/22/18:03
 * @Description:
 */
public class HeaderClientInterceptor implements ClientInterceptor {

    private static final Logger logger = Logger.getLogger(HeaderClientInterceptor.class.getName());

    @VisibleForTesting
    static final Metadata.Key<String> SERVICE_HEADER_KEY =
            Metadata.Key.of("service_client_header_key", Metadata.ASCII_STRING_MARSHALLER);

    private final String serviceName;
    public HeaderClientInterceptor(String serviceName){
        this.serviceName = serviceName;
    }

    @Override
    public <ReqT, RespT> ClientCall<ReqT, RespT> interceptCall(MethodDescriptor<ReqT, RespT> method,
                                                               CallOptions callOptions, Channel next) {
        return new ForwardingClientCall.SimpleForwardingClientCall<ReqT, RespT>(next.newCall(method, callOptions)) {

            @Override
            public void start(Listener<RespT> responseListener, Metadata headers) {
                /* put custom header */
                headers.put(SERVICE_HEADER_KEY, serviceName);
                super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(responseListener) {
                    @Override
                    public void onHeaders(Metadata headers) {
                        super.onHeaders(headers);
                    }
                }, headers);
            }
        };
    }
}
